import concurrent.futures
import requests
import os
import random
import string
import logging
from openpyxl import Workbook
from time import time, strftime, localtime

# 更新日志配置以包括线程名
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(threadName)s - %(message)s')

# 改进的参数化
num_files_to_generate = 200  # 生成文件的数量
size_in_kilobytes = 1024  # 文件大小(以KB为单位)
output_folder = 'filesOutput'  # 输出脚本所在目录文件夹
url = os.getenv('UPLOAD_URL', 'http://10.10.2.236:32287/pay/upload')  # 环境变量中读取或使用默认值
app_id = os.getenv('APP_ID', 'ndtzOjl7')
concurrent_requests = int(os.getenv('CONCURRENT_REQUESTS', 50))  # 并发数
test_duration_seconds = int(os.getenv('TEST_DURATION_SECONDS', 1))  # 持续时间（秒）

# 使用环境变量或默认值参数化请求头
headers = {
    'X-App-Id': app_id,
    'X-Timestamp': strftime('%Y%m%d%H%M%S', localtime()),
    'X-Sequence-No': '1',  # 应根据实际需求进行适当增减
    'X-Signature': 'signature'
}


def generate_files(num_files, size_in_kb, output_folder):
    """生成指定数量和大小的文件并放入指定文件夹"""
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    filenames = []
    for _ in range(num_files):
        size_in_bytes = size_in_kb * 1024  # 转换为字节
        filename = ''.join(random.choices(string.ascii_lowercase + string.digits, k=10)) + '.txt'
        filepath = os.path.join(output_folder, filename)
        with open(filepath, 'wb') as f:
            f.write(os.urandom(size_in_bytes))
        filenames.append(filepath)
    return filenames


# 修改 send_request_and_delete_file 函数以包括错误处理和响应内容处理
def send_request_and_delete_file(url, headers, filepath):
    """发送文件以及自定义请求头到url，记录响应码、内容和响应时间."""
    try:
        with open(filepath, 'rb') as f:
            filename = os.path.basename(filepath)
            files = {'file': (filename, f)}
            start_time = time()
            response = requests.post(url, headers=headers, files=files, timeout=60)  # 添加请求超时
            response_time = time() - start_time
            if response.status_code == 200:
                response_content = response.content
            else:
                response_content = response.text
            logging.info(f'{filepath} - 响应码: {response.status_code}, 响应内容: {response_content}, '
                         f'响应时间: {response_time}s')
            return filepath, response.status_code, response_time, response_content
    except requests.RequestException as e:  # 所有requests相关的异常
        logging.error(f'上传文件时出错 {filepath}: {e}')
        return filepath, None, None, str(e)


def cleanup_files(filepaths):
    """删除所有生成的文件"""
    for filepath in filepaths:
        try:
            os.remove(filepath)
        except OSError as e:
            logging.warning(f"删除文件发生错误 {filepath}: {e}")


def record_to_excel(records):
    """将记录保存为Excel文件"""
    wb = Workbook()
    ws = wb.active
    ws.append(['文件路径', '响应码', '响应时间', '响应内容'])

    success_count = 0
    if records:  # 在进行任何除法操作之前检查列表是否为空
        for record in records:
            ws.append(record)
            if record[1] == 200:  # 如果响应码为200，则将其视为成功
                success_count += 1

        # 计算成功率
        success_rate = success_count / len(records) * 100
    else:
        success_rate = 0  # 如果没有记录，则成功率为0%

    ws.append([])
    ws.append(['总请求', len(records)])
    ws.append(['成功请求', success_count])
    ws.append(['成功率', success_rate])

    wb.save("results.xlsx")


# 文件生成和测试逻辑稍有修改以限制总的执行时间
file_paths = generate_files(num_files_to_generate, size_in_kilobytes, output_folder)

start_time = time()
records = []

with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_requests) as executor:
    future_to_file = {executor.submit(send_request_and_delete_file, url, headers, fp): fp for fp in file_paths}
    try:
        for future in concurrent.futures.as_completed(future_to_file, timeout=test_duration_seconds):
            result = future.result()
            records.append(result)
            # 顺序编号对于计算多个文件的X-Sequence-No时可能有用
            headers['X-Sequence-No'] = str(int(headers['X-Sequence-No']) + 1)
    except concurrent.futures.TimeoutError:
        logging.info('测试时间已经用完，当前任务将不会等待其他任务.')

# 测试结束，记录结果到Excel
record_to_excel(records)

# 清理生成的文件
cleanup_files(file_paths)
logging.info('测试和清理结束，结果生成在 results.xlsx.')